import unittest
import sys
sys.path.append("..")

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.mcts2 import *
from psychsim.mctsWorld import *
#from psychsim.agent import Agent
from psychsim.reward import *
from psychsim.mctsAgent import mctsAgent

'''
This file serves as a template format for unit test cases.
'''

class TestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.world = mctsWorld()
        self.ralph = mctsAgent('Ralph')
        self.world.addAgent(self.ralph)
        self.sam = mctsAgent('Sam')
        self.world.addAgent(self.sam)
        self.world.setOrder(['Ralph', 'Sam'])
        
    def addStates(self):
        #define states
        self.world.defineState('Ralph','pledged', bool)
        self.world.setState('Ralph','pledged', False)
        
        self.world.defineState('Sam','asleep', bool)
        self.world.setState('Sam','asleep', False)

        self.world.defineState(None, 'sheep alert', bool)
        self.world.setState(None, 'sheep alert', False)

        self.world.defineState(None, 'sheep eaten', bool)
        self.world.setState(None, 'sheep eaten', False)

        self.world.defineState(None, 'plot thwarted', bool)
        self.world.setState(None, 'plot thwarted', False)
                   
    def addTermination(self):
        #terminate
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None, 'sheep eaten'), True),
                                            True: True,
                                            False: False}))
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None, 'plot thwarted'), True),
                                            True: True,
                                            False: False}))
        
    def addActions(self):
        #add actions and appropriate dynamics
        pledge = self.ralph.addAction({'verb': 'pledge'})
        tree = makeTree(setTrueMatrix(stateKey('Ralph', 'pledged')))
        self.world.setDynamics(stateKey('Ralph', 'pledged'), pledge, tree)
        #self.ralph.setLegal(pledge, makeTree({'if': equalRow(stateKey('Ralph', 'pledged'), False),
        #                                      True: True,
        #                                      False: False}))

        steal = self.ralph.addAction({'verb': 'steal sheep'})
        tree = makeTree(setTrueMatrix(stateKey(None, 'sheep alert')))
        self.world.setDynamics(stateKey(None, 'sheep alert'), steal, tree)
        self.ralph.setLegal(steal, makeTree({'if': equalRow(stateKey(None, 'sheep alert'), False),
                                             True: True,
                                             False: False}))

        eat = self.ralph.addAction({'verb': 'eat sheep'})
        tree = makeTree(setTrueMatrix(stateKey(None, 'sheep eaten')))
        self.world.setDynamics(stateKey(None, 'sheep eaten'), eat, tree)
        self.ralph.setLegal(eat, makeTree({'if': equalRow(stateKey(None, 'sheep alert'), True),
                                             True: True,
                                             False: False}))

        sleep = self.sam.addAction({'verb': 'sleep'})
        tree = makeTree(setTrueMatrix(stateKey('Sam', 'asleep')))
        self.world.setDynamics(stateKey('Sam', 'asleep'), sleep, tree)

        stare = self.sam.addAction({'verb': 'stare'})
        self.sam.setLegal(stare, makeTree({'if': equalRow(stateKey('Sam', 'asleep'), False),
                                             True: True,
                                             False: False}))

        thwart = self.sam.addAction({'verb': 'thwart'})
        tree = makeTree(setTrueMatrix(stateKey(None, 'plot thwarted')))
        self.world.setDynamics(stateKey(None, 'plot thwarted'), thwart, tree)
        self.sam.setLegal(thwart, makeTree({'if': equalRow(stateKey('Sam', 'asleep'), False),
                                            True: {'if': equalRow(stateKey(None, 'sheep alert'), True),
                                                   True: True,
                                                   False: False},
                                            False: False}))

        
    def setModelAndRewards(self):
        #set models, rewards, and beliefs
        self.ralph.setReward(achieveFeatureValue(stateKey(None, 'sheep eaten'), True), 10.0, True)
        
        self.ralph.addModel('Well-intentioned')
        self.ralph.setReward(achieveFeatureValue(stateKey('Sam', 'asleep'), True), 1.0, 'Well-intentioned')
        self.ralph.setReward(achieveFeatureValue(stateKey(None, 'sheep alert'), True), -1.0, 'Well-intentioned')
        self.ralph.setReward(achieveFeatureValue(stateKey('Ralph', 'pledged'), True), 1.0, 'Well-intentioned') 
        
        self.ralph.addModel('Sinister')
        self.ralph.setReward(achieveFeatureValue(stateKey(None, 'sheep alert'), True), 1.0, 'Sinister')
        self.ralph.setReward(achieveFeatureValue(stateKey(None, 'sheep eaten'), True), 10.0, 'Sinister')

        self.sam.setReward(achieveFeatureValue(stateKey(None, 'plot thwarted'), True), 10.0, True)
        self.sam.setReward(achieveFeatureValue(stateKey('Sam', 'asleep'), True), 1.0, True)

        self.world.setMentalModel('Sam', 'Ralph', {'Well-intentioned': .01, 'Sinister': .99})
        self.world.setModel('Sam', True)

    def setAttributes(self):
        self.ralph.setAttribute('discount', .9)
        self.ralph.setAttribute('horizon',6)
        self.sam.setAttribute('discount', .9)
        self.sam.setAttribute('horizon',6)
    
    def testSetup(self):
        #test scenario construction methods
        self.addStates()
        self.addTermination()
        self.addActions()
        self.setModelAndRewards()
        self.setAttributes()

    def testRun(self):
        #set up the scenario
        self.testSetup()
        #test running the scenario
        self.world.printState(beliefs=True)
        self.world.printModels()

        stats = getResultVariance(self.world, 10)
        print stats
        count = 0
        actionList = []
        while not self.world.terminated():
            action = psychSimUTC(self.world, 10, horizon = 3, stat = stats)
            print action
            actionList.append(action)
            self.world.DoMove(action)
            #self.world.explain(self.world.step(),level=2)
            self.world.printState(beliefs=True)
            self.world.printModels()
            #self.world.printModels()
            count += 1
            if count == 50:
                break
        for action in actionList:
            print action

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(TestCase):
       if method.startswith("test"):
          suite.addTest(TestCase(method))
    unittest.TextTestRunner().run(suite)
            
